##############################################################################
# File-Name: _utils.r
# Purpose: auxiliary functions
# Machine: macOS Ventura 13.5.1
# R version 4.3.1 
##############################################################################


# Packages ----------------------------------------------------------------
pacman::p_load(tidyverse, scales,  knitr, srvyr, extrafont, 
               rebus, broom, tidyr, lubridate, here, ggdist, ggtext, ggalt,  janitor, 
               broom, tidyr,  stargazer, janitor, estimatr, list)


# Functions ---------------------------------------------------------------

# means difference
func_t.test = function(depvar, data, model=FALSE){
  
  # separate treatment and control
  control <- data %>% filter(exp=="control") %>% pull({{depvar}})
  treatment <- data %>% filter(exp=="treatment") %>% pull({{depvar}})
  
  # t.test
  out = t.test(treatment, control)
  names(out$estimate) <- c("mean of treatment", "mean of control")
  return(out)
   
}

# ITT: ols
# inputs: dependent variable, data, and logical for model, where false returns the full model, and true returns a dataframe with only the standardized treatment effects. 
func_ols_base = function(depvar, data, model=FALSE){

if(model==TRUE){
model.expr = as.formula(paste(depvar, " ~ exp"))
out = lm_robust(model.expr, data = data, se_type = "stata")
return(out)}
  
model.expr = as.formula(paste(depvar, " ~ exp"))
out = lm_robust(model.expr, data = data, se_type = "HC0")
  
  # sd control
  sd_control <- data %>% 
    filter(exp=="control") %>%
    pull(depvar) %>%
    sd(., na.rm=TRUE)
  
  
  # tidy
  out <- out %>% 
    tidy() %>%
    filter(term=="exptreatment") %>%
    mutate_at(vars(c(estimate, std.error)), 
              ~.x/sd_control)
  return(out)
  
}

# Covariate adjusted ITT: ols + controls
# inputs: dependent variable, data, and logical for model, where false returns the full model, and true returns a dataframe with only the standardized treatment effects. 
func_ols_covariates = function(depvar, data, model=FALSE){
  # For the controls, we will use
  # respondents’ demographics (sex, age, education, income, race), measures of trust in mainstream media
  # and institutions, pre-treatment measures of polarization, ideology and social media usage.
  add.vars = c(           
    #demographics
      "w1_q_gender" ,  
      "q_race" ,  
      "w1_q_education" ,
      "w1_q_age_num",
      "income_num",

      
      
      # trust  
      "as.numeric(w1_q_trust_government)" ,
      "as.numeric(w1_q_trust_congress)" ,
      "as.numeric(w1_q_trust_supreme_court)" ,
      "as.numeric(w1_q_trust_electoral_authorities)" ,
      "as.numeric(w1_q_trust_globo)" ,
      "as.numeric(w1_q_trust_news_channels)" ,
      "w1_q_legitimacy" ,
      
      # polarization    
      "w1_affective_pol_bolsonaro" ,
       "w1_affective_pol_lula",

      # ideology  
      "w1_ideology_you" ,  
      "w1_q_politics_num" ,   
      
      # social media usage
      "w1_q_whatsapp_num ",
      "as.numeric(w1_q_fake_news)" ,
      "as.numeric(w1_q_whatsapp_purposes_news)" ,
      "as.numeric(w1_q_whatsapp_purposes_pay_bills)" ,
      "as.numeric(w1_q_whatsapp_purposes_family)" ,
      "as.numeric(w1_q_whatsapp_purposes_groups)"
     )

  if(model==TRUE){
    model.expr = as.formula(paste(depvar, " ~ exp +",
                                  paste(add.vars, collapse = "+")))
    
    out = lm_robust(model.expr, data = data, se_type="HC0")
    return(out)
  }    
  
  model.expr = as.formula(paste(depvar, " ~ exp +",
                                paste(add.vars, collapse = "+")))
  
  out = lm_robust(model.expr, data = data, se_type = "HC0")
  # sd control
  sd_control <- data %>% 
    filter(exp=="control") %>%
    pull(depvar) %>%
    sd(., na.rm=TRUE)
  
  # tidy
  out <- out %>% 
    tidy() %>%
    filter(term=="exptreatment") %>%
    mutate_at(vars(c(estimate, std.error)), 
              ~.x/sd_control)
  
  return(out)
}


# CACE
# inputs: dependent variable, data, and logical for model, where false returns the full model, and true returns a dataframe with only the standardized treatment effects. 
func_cace = function(depvar, data, model=FALSE){
  
if(model==TRUE){
  model.expr = as.formula(paste(depvar, " ~ compliance_one_side | exp "))
  
  # model
  out = iv_robust(model.expr, data = data) 
  return(out)
}
  model.expr = as.formula(paste(depvar, " ~ compliance_one_side | exp "))
  
  # model
  out = iv_robust(model.expr, data = data) 
  
  # sd control
  sd_control <- data %>% 
    filter(exp=="control") %>%
    pull(depvar) %>%
    sd(., na.rm=TRUE)
  
  # tidy
  out <- out %>% 
    tidy() %>%
    mutate_at(vars(c(estimate, std.error)), 
              ~.x/sd_control)
  return(out)
}


# ggplot theme for the paper ------------------------------------------------------------
my_font <- "Verdana"
my_bkgd <- "white"
#my_bkgd <- "#f5f5f2"
pal <- RColorBrewer::brewer.pal(9, "Spectral")
colorwes <-  wesanderson::wes_palette("BottleRocket2")[[3]]

my_theme <- theme(text = element_text(family = my_font, color = "#22211d"),
                  rect = element_rect(fill = my_bkgd),
                  plot.background = element_rect(fill = my_bkgd),
                  panel.background = element_rect(fill = my_bkgd),
                  panel.border = element_rect(color="transparent"), 
                  strip.background = element_rect(color="transparent", fill="transparent"), 
                  legend.background = element_rect(fill = my_bkgd, color = NA),
                  legend.key = element_rect(size = 4, fill = "white", colour = NA), 
                  legend.key.size = unit(1, "cm"),
                  legend.text = element_text(size = 12, family = my_font),
                  legend.title = element_text(size=16, family=my_font, face="bold"),
                  plot.title = element_markdown(size = 22, face = "bold", family=my_font),
                  plot.subtitle = element_markdown(size=16, family=my_font),
                  axis.title= element_text(size=22),
                  axis.text = element_text(size=12, family=my_font),
                  axis.title.x = element_text(hjust=1),
                  strip.text = element_text(family = my_font, 
                                            color = "#22211d",
                                            size = 14, face="bold"), 
                  plot.margin = margin(1, 1, 1, 1, "cm"), 
                  plot.caption = element_text(size=12, 
                                              family=my_font,
                                              hjust = 0.5, 
                                              face="italic"))


theme_set(theme_bw() + my_theme)

## ggpoint
plot_model <- function(output, 
                       models,
                       dep_var_label,
                       colorvalues=colorwes,
                       data,
                       xlab="", ylab="", 
                       title="", subtitle="", caption="", style=1){
  
  # first tidy
  tidydata <- pmap(list(output, 
                        models, 
                        dep_var_label), 
                   function(data, model, dv) 
                     data %>%
                     mutate(model=model) %>%
                     filter(term=="exptreatment"|term=="compliance_one_side") %>%
                     mutate(up90=estimate + 1.64*std.error,
                            up95=estimate + 1.96*std.error, 
                            lb90=estimate - 1.64*std.error, 
                            lb95=estimate - 1.96*std.error) %>%
                     mutate(dv=dv)) %>%
    bind_rows()  %>%
    mutate(model=fct_relevel(as.factor(model), c("CACE", "Cov-ITT", "ITT")))
  
  
  
  
  
  # graph
  if(style==1){
    
    # graph with a unique DV
    ggplot(tidydata,
           aes(y=model, 
               x=estimate, 
               color=model,
               label=round(estimate, 2))) +
      geom_errorbar(aes(xmin=lb90, 
                        xmax=up90), 
                    size=2, width=0, alpha=.8, 
                    position=position_dodge(width = .8), 
                    color=colorvalues) +
      geom_errorbar(aes(xmin=lb95, 
                        xmax=up95),  
                    size=1, width=.1, alpha=.5, 
                    position=position_dodge(width = .8), 
                    color=colorvalues) +
      geom_point(size=14, shape=21, fill="white",
                 position=position_dodge(width = .6), 
                 stroke=2, 
                 color=colorvalues, alpha=1) +
      geom_text(
        fontface = "italic", 
        fill="white", 
        color = "black",
        size=4, 
        position = position_dodge(.6)) +
      geom_vline(aes(xintercept=0), linetype="dashed", color="tomato2", alpha=.8) +
      ylab(ylab) +
      xlab(xlab) +
      labs(title=title, 
           subtitle=subtitle, 
           caption=caption) +
      # scale_color_manual(values=colorvalues,  name="Models", 
      #                   guide = guide_legend(reverse = TRUE, 
      #                                        override.aes = list(size=3), 
      #                                       title.position="left")) +
      facet_grid(~dv) +
      theme(legend.position = "none")
  }
  else{
    # graph with a unique DV
    ggplot(tidydata,
           aes(y=dv, 
               x=estimate, 
               color=model)) +
      geom_errorbar(aes(xmin=lb90, 
                        xmax=up90), 
                    size=2, width=0, alpha=1, 
                    position=position_dodge(width = .2)) +
      geom_errorbar(aes(xmin=lb95, 
                        xmax=up95),  
                    size=1, width=.2, alpha=.7, 
                    position=position_dodge(width = .2)) +
      geom_point(size=2, shape=22, fill="white",
                 position=position_dodge(width = 0.2), alpha=.8) +
      geom_vline(aes(xintercept=0), linetype="dashed", color="tomato2") +
      ylab(ylab) +
      xlab(xlab) +
      labs(title=title, 
           subtitle=subtitle) +
      scale_color_manual(values=colorvalues,  name="Models", 
                         guide = guide_legend(reverse = TRUE))  
    
  }
}


# function for interactive marginal effect with a binary moderator
interaction_plot_binary <- function(model, effect, moderator, interaction, varcov="default", conf=.95, title="Marginal effects plot", xlabel="Value of moderator", ylabel="Estimated marginal coefficient", factor_labels=c(0,1)){
  
  # Extract Variance Covariance matrix
  if (varcov == "default"){
    covMat = vcov(model)
  }else{
    covMat = varcov
  }
  
  # Extract the data frame of the model
  mod_frame = model.frame(model)
  
  # Get coefficients of variables
  beta_1 = model$coefficients[[effect]]
  beta_3 = model$coefficients[[interaction]]
  
  # Create list of moderator values at which marginal effect is evaluated
  x_2 <- c(0,1)
  
  # Compute marginal effects
  delta_1 = beta_1 + beta_3*x_2
  
  # Compute variances
  var_1 = covMat[effect,effect] + (x_2^2)*covMat[interaction, interaction] + 2*x_2*covMat[effect, interaction]
  
  # Standard errors
  se_1 = sqrt(var_1)
  
  # Upper and lower confidence bounds
  z_score = qnorm(1 - ((1 - conf)/2))
  upper_bound = delta_1 + z_score*se_1
  lower_bound = delta_1 - z_score*se_1
  
  # Determine the bounds of the graphing area
  max_y = max(upper_bound)
  min_y = min(lower_bound)
  
  # Initialize plotting window
  plot(x=c(), y=c(), ylim=c(min_y, max_y), xlim=c(-.5, 1.5), xlab=xlabel, ylab=ylabel, main=title, xaxt="n")
  
  # Plot points of estimated effects
  points(x=x_2, y=delta_1, pch=16)
  
  # Plot lines of confidence intervals
  lines(x=c(x_2[1], x_2[1]), y=c(upper_bound[1], lower_bound[1]), lty=1)
  points(x=c(x_2[1], x_2[1]), y=c(upper_bound[1], lower_bound[1]), pch=c(25,24), bg="black")
  lines(x=c(x_2[2], x_2[2]), y=c(upper_bound[2], lower_bound[2]), lty=1)
  points(x=c(x_2[2], x_2[2]), y=c(upper_bound[2], lower_bound[2]), pch=c(25,24), bg="black")
  
  # Label the axis
  axis(side=1, at=c(0,1), labels=factor_labels)
  
  # Add a dashed horizontal line for zero
  abline(h=0, lty=3)
  
  out = tibble(delta_1, upper_bound, lower_bound, moderator=x_2, treatment=effect, std.error=se_1)
  
  
}


# Functiont to save graph outputs locally and on overleaf ----------------------------------
save_function <- function(name){
  
  #overleaf output
  #output_ov <- "~/Dropbox/Apps/Overleaf/whatsapp_deactivation/output"
  
  ggsave(filename = here("output", name), 
         width = 12, height = 8, units = "in", 
         pointsize = 12, bg = "white")
  
 # ggsave(filename = paste0(output_ov, "/", name), 
  #       width = 12, height = 8, units = "in", 
   #      pointsize = 12, bg = "white") 

}

# function to get marginal effects for continuous moderator
interaction_plot_continuous_ggplot <- function(model, effect, 
                                               moderator,
                                               interaction, 
                                               varcov="default", 
                                               minimum="min", 
                                               maximum="max", 
                                               incr="default", 
                                               num_points = 50, conf=.95, mean=FALSE, median=FALSE,
                                               alph=80, rugplot=T, histogram=T, title="Marginal effects plot",
                                               xlabel="Value of moderator", ylabel="Estimated marginal coefficient"){
  
  
  # Extract Variance Covariance matrix
  if (varcov == "default"){
    covMat = vcov(model)
  }else{
    covMat = varcov
  }
  
  # Extract the data frame of the model
  mod_frame = model.frame(model)
  
  # Get coefficients of variables
  beta_1 = model$coefficients[[effect]]
  beta_3 = model$coefficients[[interaction]]
  
  # Set range of the moderator variable
  # Minimum
  if (minimum == "min"){
    min_val = min(mod_frame[[moderator]])
  }else{
    min_val = minimum
  }
  # Maximum
  if (maximum == "max"){
    max_val = max(mod_frame[[moderator]])
  }else{
    max_val = maximum
  }
  
  # Check if minimum smaller than maximum
  if (min_val > max_val){
    stop("Error: Minimum moderator value greater than maximum value.")
  }
  
  # Determine intervals between values of the moderator
  if (incr == "default"){
    increment = (max_val - min_val)/(num_points - 1)
  }else{
    increment = incr
  }
  
  # Create list of moderator values at which marginal effect is evaluated
  x_2 <- seq(from=min_val, to=max_val, by=increment)
  
  # Compute marginal effects
  delta_1 = beta_1 + beta_3*x_2
  
  # Compute variances
  var_1 = covMat[effect,effect] + (x_2^2)*covMat[interaction, interaction] + 2*x_2*covMat[effect, interaction]
  
  # Standard errors
  se_1 = sqrt(var_1)
  
  # Upper and lower confidence bounds
  z_score = qnorm(1 - ((1 - conf)/2))
  upper_bound = delta_1 + z_score*se_1
  lower_bound = delta_1 - z_score*se_1
  
  # Determine the bounds of the graphing area
  max_y = max(upper_bound)
  min_y = min(lower_bound)
  
  # Make the histogram color
  
  # Initialize plotting window
  plot(x=c(), y=c(), ylim=c(min_y, max_y), xlim=c(min_val, max_val), xlab=xlabel, ylab=ylabel, main=title)
  
  # Plot estimated effects
  lines(y=delta_1, x=x_2)
  lines(y=upper_bound, x=x_2, lty=2)
  lines(y=lower_bound, x=x_2, lty=2)
  
  # Add a dashed horizontal line for zero
  abline(h=0, lty=3)
  
  # # Add a vertical line at the mean
  # if (mean){
  #   abline(v = mean(mod_frame[[moderator]]), lty=2, col="red")
  # }
  # 
  # # Add a vertical line at the median
  # if (median){
  #   abline(v = median(mod_frame[[moderator]]), lty=3, col="blue")
  # }
  # 
  # # Add Rug plot
  # if (rugplot){
  #   rug(mod_frame[[moderator]])
  # }
  # 
  out = tibble(delta_1, upper_bound, lower_bound, moderator=x_2, treatment=effect)
  
  return(out)
}

# fucntion to get marginal effects binary moderator
interaction_plot_binary <- function(model, effect, moderator, interaction, 
                                    varcov="default", conf=.95, title="Marginal effects plot", 
                                    xlabel="Value of moderator", ylabel="Estimated marginal coefficient", factor_labels=c(0,1)){
  
  # Extract Variance Covariance matrix
  if (varcov == "default"){
    covMat = vcov(model)
  }else{
    covMat = varcov
  }
  
  # Extract the data frame of the model
  #mod_frame = model.frame(model)
  
  # Get coefficients of variables
  beta_1 = model$coefficients[[effect]]
  beta_3 = model$coefficients[[interaction]]
  
  # Create list of moderator values at which marginal effect is evaluated
  x_2 <- c(0,1)
  
  # Compute marginal effects
  delta_1 = beta_1 + beta_3*x_2
  
  # Compute variances
  var_1 = covMat[effect,effect] + (x_2^2)*covMat[interaction, interaction] + 2*x_2*covMat[effect, interaction]
  
  # Standard errors
  se_1 = sqrt(var_1)
  
  # Upper and lower confidence bounds
  z_score = qnorm(1 - ((1 - conf)/2))
  upper_bound = delta_1 + z_score*se_1
  lower_bound = delta_1 - z_score*se_1
  
  # Determine the bounds of the graphing area
  max_y = max(upper_bound)
  min_y = min(lower_bound)
  
  # Initialize plotting window
  plot(x=c(), y=c(), ylim=c(min_y, max_y), xlim=c(-.5, 1.5), xlab=xlabel, ylab=ylabel, main=title, xaxt="n")
  
  # Plot points of estimated effects
  points(x=x_2, y=delta_1, pch=16)
  
  # Plot lines of confidence intervals
  lines(x=c(x_2[1], x_2[1]), y=c(upper_bound[1], lower_bound[1]), lty=1)
  points(x=c(x_2[1], x_2[1]), y=c(upper_bound[1], lower_bound[1]), pch=c(25,24), bg="black")
  lines(x=c(x_2[2], x_2[2]), y=c(upper_bound[2], lower_bound[2]), lty=1)
  points(x=c(x_2[2], x_2[2]), y=c(upper_bound[2], lower_bound[2]), pch=c(25,24), bg="black")
  
  # Label the axis
  axis(side=1, at=c(0,1), labels=factor_labels)
  
  # Add a dashed horizontal line for zero
  abline(h=0, lty=3)
  
  out = tibble(delta_1, std.error=se_1, upper_bound, lower_bound, moderator=x_2, treatment=effect)
  
  
}


